#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''=================================================
@Project -> File   ：db-toolkit -> database
@IDE    ：PyCharm
@Author ：Json
@Date   ：2025/3/24 11:55
@Desc   ：
=================================================='''
# database.py
from sqlalchemy import create_engine, inspect, text
from sqlalchemy.engine import Engine
from sqlalchemy.exc import OperationalError, SQLAlchemyError
import functools
import time
from typing import List, Dict, Union
import pandas as pd
from .config import DBConfig
from .singleton import singleton


def _auto_retry(func):
    """自动重试装饰器（最大3次）"""

    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        retries = 3
        for attempt in range(retries):
            try:
                return func(self, *args, **kwargs)
            except OperationalError as e:
                if attempt < retries - 1 and e.orig.args[0] in (2006, 2013):
                    time.sleep(2 ** attempt)  # 指数退避
                    self.reconnect()
                    continue
                raise
        return None

    return wrapper


@singleton
class DBS:
    def __init__(self, database_name='amazon', config_path=None):
        self._engine = None
        self.config = DBConfig.load(config_path, database_name)

    @property
    def engine(self) -> Engine:
        if not self._engine:
            self._engine = self._create_engine()
        return self._engine

    def _create_engine(self) -> Engine:
        try:
            return create_engine(
                f"mysql+pymysql://{self.config['user']}:{self.config['password']}"
                f"@{self.config['host']}:{self.config['port']}/{self.config['database']}"
                "?charset=utf8mb4&connect_timeout=10",
                pool_size=5,
                max_overflow=10,
                pool_recycle=3600
            )
        except Exception as e:
            raise RuntimeError(f"数据库连接失败: {str(e)}")

    def reconnect(self):
        """显式重连"""
        if self._engine:
            self._engine.dispose()
        self._engine = self._create_engine()

    @_auto_retry
    def query(self, sql):  # 执行原生sql查询语句
        """兼容以前的"""
        return self.execute(sql)

    def operation(self, sql):
        """兼容以前的"""
        self.execute(sql)

    def execute(self, sql: str, params: Union[Dict, List] = None):
        """执行查询语句"""
        with self.engine.connect() as conn:
            result = conn.execute(text(sql), params) if params else conn.execute(text(sql))
            if result.returns_rows:
                return result.fetchall()
            else:
                return None

    @_auto_retry
    def execute_many(self, sql: str, params_list: List[Dict]):
        """批量执行"""
        with self.engine.begin() as conn:
            conn.execute(text(sql), params_list)

    def save_multiple_data(self, sql, data_list):
        with self.engine.connect() as con:
            con.execute(sql, data_list)
            return True

    def save_multiple_data_df(self, table_name, data_df, insert_type='INSERT'):
        with self.engine.connect() as con:
            # columns_name = data_df.columns  # dataframe的列名，是数据库的列名的子集
            # columns_name_str = ','.join(columns_name)
            # style_str = ','.join(['%s'] * len(columns_name))

            columns = data_df.columns.tolist()
            columns_name_str = ','.join(columns)
            style_str = ','.join([f":{col}" for col in columns])
            sql = f"{insert_type} INTO {table_name} ({columns_name_str}) VALUES ({style_str});"
            # 将DataFrame转换为字典列表
            data_list = data_df.to_dict(orient='records')
            print(sql)
            con.execute(text(sql), data_list)
            return True

    def save_dict_data(self, table_name, dict_data, insert_type='INSERT'):
        with self.engine.connect() as con:
            columns_names = ', '.join(dict_data.keys())
            style_str = ','.join(['%s'] * len(dict_data.keys()))
            # 使用 INSERT INTO 语法
            sql = f"{insert_type} INTO {table_name} ({columns_names}) VALUES ({style_str})"
            con.execute(text(sql), [dict_data.values()])  # 使用参数绑定传递字典数

    def read_sql_query(self, sql):
        import pandas as pd
        result = pd.read_sql_query(sql, self.engine)
        return result

    def get_columns(self, table: str):
        """获取表字段"""
        insp = inspect(self.engine)
        return [col['name'] for col in insp.get_columns(table)]

    def fitter_columns(self, data_df_columns, table_name):
        result = self.get_columns(table_name)
        columns_set = {item['name'] for item in result}
        fitter_columns = set(data_df_columns) & columns_set
        return list(fitter_columns)

    def health_check(self) -> bool:
        """健康检查"""
        try:
            self.execute("SELECT 1")
            return True
        except SQLAlchemyError:
            return False
